import pickle

import numpy as np

prefix = './'


def _load_batch(fpath, label_key=b'labels'):
    """Internal utility for parsing CIFAR data.
  
    Arguments:
        fpath: path the file to parse.
        label_key: key for label data in the retrieve
            dictionary.
  
    Returns:
        A tuple `(data, labels)`.
    """
    _dic = _unpickle(fpath)
    _data = _dic[b'data']
    _labels = _dic[label_key]
    
    _data = _data.reshape(_data.shape[0], 3, 32, 32)
    return _data, _labels


def _unpickle(file):
    """
    unpickle files
    :param file:
    :return:
    """
    with open(file, 'rb') as fo:
        _dic = pickle.load(fo, encoding='bytes')
    return _dic


def train_data():
    """
    return all train data and label with numpy array
    :return:
    """
    num_train_samples = 50000
    x_train = np.empty((num_train_samples, 3, 32, 32), dtype='uint8')
    y_train = np.empty((num_train_samples,), dtype='uint8')
    
    for i in range(1, 6):
        fpath = prefix + 'cifar-10-batches-py/data_batch_' + str(i)
        (x_train[(i - 1) * 10000:i * 10000, :, :, :],
         y_train[(i - 1) * 10000:i * 10000]) = _load_batch(fpath)
    
    # 'channels_first' or 'channels_last'
    # 转换图片的格式，符合框架的输入要求
    x_train = np.transpose(x_train, (0, 2, 3, 1))
    return x_train, y_train


def test_data():
    """
    return all test data and label with numpy array
    :return:
    """
    x_test, y_test = _load_batch(prefix + 'cifar-10-batches-py/test_batch')

    # 'channels_first'` or `'channels_last'
    # 转换图片的格式，符合框架的输入要求
    x_test = np.transpose(x_test, (0, 2, 3, 1))
    return x_test, y_test


def meta_data():
    """
    return meta data
    :return:
    """
    return _unpickle(prefix + 'cifar-10-batches-py/batches.meta')


def categories():
    return meta_data()[b'label_names']


if __name__ == '__main__':
    pass

